import os
import subprocess
import fnmatch
import numpy as np

class simulation:

	def __init__(self,dir):
		files = fnmatch.filter(os.listdir(dir),'profiles_*')
		self.iterations = []
		self.times = []
		for file in files:
			(self.iterations).append(file.split('_')[1])
			f = open(file,'r')
			header = f.readline()
			header = header.split()
			(self.times).append(header[4])
			self.aspect_ratio = float(header[-1])
			self.nx = int(header[6])
			self.nz = int(header[8])

		print "   "+(self.iterations)[-1]+" iterations"
		print "   %06d outputs" % len(files)

	def getGrids(self, iter):
		data = {}
		grids = "grd_"+str(iter)+"_.dat"
		[data['x'], data['z'], data['vx'], data['vz'], data['velocity'], data['temperature'], data['viscosity'], data['visc_diss'], \
				data['composition'], data['RaH'], data['expansivity'], data['conductivity']] = np.loadtxt(grids, unpack=True, skiprows=1)
	
		data['x'] = data['x'][::self.nz]
		data['z'] = data['z'][:self.nz]
		data['z'] = data['z'][::-1]


		data['vx'] = data['vx'].reshape(self.nx,self.nz)
		data['vx'] = data['vx'].T
		data['vx'] = data['vx'][::-1,:]
			
		data['vz'] = data['vz'].reshape(self.nx,self.nz)
		data['vz'] = data['vz'].T
		data['vz'] = data['vz'][::-1,:]

		return data

	def getProfiles(self, iter):
		data = {}
		profiles = "profiles_"+str(iter)+"_.dat"
		[data['depth'], data['velocity'], data['temperature'], data['viscosity'], data['visc_diss'], data['composition'], \
				data['vx'], data['vz'], data['RaH'], data['expansivity'], data['conductivity']] = np.loadtxt(profiles, unpack=True, skiprows=1)
		return data

	def getTimeseries(self):
		data = {}
		filename = "timeseries.dat"
		[data['time'], data['nu'], data['temperature'], data['velocity'], data['vsurf'], data['viscosity'], \
				data['visc_diss'], data['composition']] = np.loadtxt(filename, unpack=True, skiprows=1)
		return data
